"""
这个文件用来将数据的bounddingbox裁剪成一张张图片，然后创建对应数据集
一个.jpg，一个.txt  -------------->     n张.jpg
"""
#images_ave_path
#   dog
#       1.jpg
#       2.jpg
#       3.jpg
#   cat
#       7.jpg
#       1.jpg
from tqdm import tqdm
import cv2
import os

origin_images_path = "/media/jiji/fe0e60a9-bc54-4761-a52c-dd4ec10ba6db/Company/我的数据集/来源/obj_det_plate/plate_train_data/VOC2007/JPEGImages/"     #原始图像
origin_label_path = "/media/jiji/fe0e60a9-bc54-4761-a52c-dd4ec10ba6db/Company/我的数据集/来源/obj_det_plate/plate_train_data/VOC2007/labels/"          #原始标签
images_save_path = "/media/jiji/fe0e60a9-bc54-4761-a52c-dd4ec10ba6db/Company/我的数据集/来源/obj_det_plate/plate_train_data/VOC2007/截取的文件夹/"       #图片文件保存的目录
classes_name_path = "/media/jiji/fe0e60a9-bc54-4761-a52c-dd4ec10ba6db/Company/我的数据集/来源/obj_det_plate/obj.name"          #存放类别名的文件夹

def read_class_names(path:str)->list:
    f = open(path,'r')
    names = []
    for line in f.readlines():
        names.append(line.strip())
    return names

def creat_directory():
    '''
    为所有的类别创建文件夹
    '''
    for name in read_class_names(classes_name_path):
        if not os.path.isdir(images_save_path+name):
            os.mkdir(images_save_path+name)

if __name__ == '__main__':
    creat_directory()
    file_names = os.listdir(origin_label_path)
    classes_name = read_class_names(classes_name_path)
    for file_name in tqdm(file_names, ncols=150):
        img = cv2.imread(origin_images_path + file_name[:-4] + ".jpg")
        img_h, img_w = img.shape[0], img.shape[1]
        txt_file = open(origin_label_path + file_name, 'r')
        for i,line in enumerate(txt_file.readlines()):
            line_chars = line.split(' ')
            class_id = int(line_chars[0])
            # 转换为float
            coordinate_x, coordianete_y = float(line_chars[1]), float(line_chars[2])
            bbox_w, bbox_h = float(line_chars[3]), float(line_chars[4])
            coordinate_x, coordianete_y = round((coordinate_x - bbox_w / 2) * img_w), round((coordianete_y - bbox_h / 2) * img_h)
            bbox_w, bbox_h = round(bbox_w * img_w), round(bbox_h * img_h)
            point_1, point_2 = [coordinate_x, coordianete_y], [coordinate_x + bbox_w, coordianete_y + bbox_h]
            croped_img = img[point_1[1]:point_2[1],point_1[0]:point_2[0]]
            if croped_img.shape[0] and croped_img.shape[1]:
                cv2.imwrite(images_save_path +str(classes_name[class_id])+os.sep+file_name[:-4] + f"{str(i)}.jpg", croped_img)
